#include "core/main.h"

Game* Application::GameFromName(const string &name)
{
    if(name == "snake") return new GameSnake();
    if(name == "snakeBasic") return new GameSnakeBasic();
    if(name == "snakeOnline") return new GameSnakeOnline();
    if(name == "testBasic") return new GameTestBasic();
    if(name == "missileSwarm") return new GameMissileSwarm();
    if(name == "breakout") return new GameBreakout();
    if(name == "gridWorld") return new GameGridWorld();
    if(name == "fruitEater") return new GameFruitEater();
    if(name == "frogger") return new GameFrogger();
    if(name == "dodgeTheMissile") return new GameDodgeTheMissile();
    if(name == "pong") return new GamePong();
    if(name == "ddr") return new GameDDR();
    Assert(false, "Game not found");
    return NULL;
}

void Application::InitParameters()
{
    _parameters.Init("../parameters.txt");
}

void Application::Init()
{
    srand((UINT)time(NULL));

    _frameBufferA.Allocate(Constants::ScreenTextureWidth, Constants::ScreenTextureHeight);
    _frameBufferB.Allocate(Constants::ScreenTextureWidth, Constants::ScreenTextureHeight);
    _frameTextureA = new GLTexture;
    _frameTextureB = new GLTexture;
    _controller.state = 0;

    _activeGame = GameFromName(_parameters.game);
    _activeGame->Init();
    _activeGame->NewGame();

    if(_parameters.watchingReplays)
    {
        _allReplays.LoadAllReplays(*_activeGame, "", _parameters.testingModel);
        _replayIndex = 0;
        _replayFrameIndex = 0;

        if(_parameters.learningModel)
        {
            ModelLearner learner;
            GameModel model;
            GameModelParameters modelParameters;

            modelParameters.gridScale = Constants::NESTileSize;
            modelParameters.historyFrameCount = 1;
            modelParameters.skipIntegers.resize(0);
            modelParameters.stateWindowSize = 1;
            modelParameters.maxAbsoluteVariableCount = 3;
            modelParameters.spriteDatasetSize = 50000 * 3;
            //modelParameters.spriteDatasetSize = 5 * 3;

            learner.LearnModel(_allReplays, modelParameters, model);
            model.Save("../models/" + _activeGame->Name() + "/model_000");
        }
        if(_parameters.testingModel)
        {
            _activeModel.Load(*_activeGame, "../models/" + _activeGame->Name() + "/model_000");
            cout << "Loaded model" << endl;
        }
        if(_parameters.learningNeighborhoods)
        {
            NeighborhoodFeatureSet featureSet;
            featureSet.buildFromReplays(_allReplays, 500, 3, Constants::NESTileSize);
            featureSet.saveToFile("../neighborhoods/" + _activeGame->Name() + "/neighborhoods.neigh");
            featureSet.exportPNGs(_activeGame->Collection(), "../neighborhoods/" + _activeGame->Name() + "/pngs/");
        }
        if(_parameters.learningValueFn)
        {
            // This requires that a featureset and a model have previously been learned for this game.
            /*
            NeighborhoodFeatureSet featureSet;
            featureSet.loadFromFile("../neighborhoods/" + _activeGame->Name() + "/neighborhoods.neigh");
            FeatureTransform transform(Constants::NESTileSize, featureSet);
            */

            TcFeatureTransformParams transformParams;
            transformParams.absPosGridSize = 2 * Constants::NESTileSize;
            transformParams.frameHeight = Constants::NESHeight;
            transformParams.frameWidth = Constants::NESWidth;
            transformParams.numSprites = _activeGame->Collection().FontBase();
            TcFeatureTransform transform(transformParams);

            _activeModel.Load(*_activeGame, "../models/" + _activeGame->Name() + "/model_000");

            ValueIterationParams viParams;
            viParams.discountFactor = 0.9;
            viParams.maxIterations = 100;
            ValueIterationLearner viLearner(viParams, transform);

            ValueFunction valueFn;
            viLearner.learnValueFn(_allReplays, _activeModel, valueFn);
            valueFn.saveToFile("../valueFns/" + _activeGame->Name() + "/valueFn_000");
        }
        if(_parameters.evaluatingModel) 
        {
            // This requires that a model has been trained.
            _activeModel.Load(*_activeGame, "../models/" + _activeGame->Name() + "/model_000");

            GameModelEvaluatorParams evaluatorParams;
            evaluatorParams.gridScale = Constants::NESTileSize;
            evaluatorParams.maxLookAhead = 20;
            stringstream dirName;
            dirName << "../confusions/" << _activeGame->Name() << "/";
            evaluatorParams.dirName = dirName.str();
            GameModelEvaluator evaluator;
            evaluator.evaluateModel(*_activeGame, _allReplays, _activeModel, evaluatorParams);
        }
    }
    else if(_parameters.recording)
    {
        const string &playerName = "AI0";
        _replay.InitRecording(*_activeGame, playerName);
    }

    if (_parameters.learningOnline)
    {
        vector<ControllerState> actions;

        if (_parameters.useFourDirectionActionSpace)
        {
            actions = ControllerState::getDirectionStates();
        }
        else
        {
            actions = ControllerState::getLeftRightStates();
        }
        
        // Initialize the transform to be used
        TcFeatureTransformParams transformParams;
        transformParams.absPosGridSize = Constants::NESTileSize;
        transformParams.relPosGridSize = Constants::NESTileSize;
        transformParams.frameHeight = Constants::NESHeight;
        transformParams.frameWidth = Constants::NESWidth;
        transformParams.numSprites = _activeGame->Collection().FontBase();
        transformParams.numActions = actions.size();
        TcFeatureTransform transform(transformParams);

        GameAgent* gameAgent = NULL;
        if (_parameters.onlineGameAgent == "sarsa")
        {
            gameAgent = new SarsaLearner(_activeGame, transform, actions);
        }
        else if (_parameters.onlineGameAgent == "alwaysUp")
        {
            ControllerState up;
            up.KeyDown('w');

            gameAgent = new AlwaysPerformSingleActionGameAgent(_activeGame, up);
        }
        else if (_parameters.onlineGameAgent == "random")
        {
            gameAgent = new RandomGameAgent(_activeGame, actions);
        }
        else
        {
            RELEASE_ASSERT( 0 );
        }

        _onlineLearningCoordinator = new SarsaCoordinator(gameAgent, _parameters.onlineLearningEpisodes);
    }
    
    _useAI = _parameters.useAI;

    _movieFrame = 0;

    _testingBaseFrame = 0;
    _testingLookAhead = 1;
    _predictedFrames.resize(_testingLookAhead);
}

void Application::Draw()
{
    if(_parameters.testingModel)
    {
        _allReplays.Replays()[0]->Frames()[_testingBaseFrame + _testingLookAhead].Render(_frameBufferA, _activeGame->Collection());
        if(_testingLookAhead != 0) _predictedFrames[_testingLookAhead - 1].Render(_frameBufferB, _activeGame->Collection());
    }
    else if(_parameters.watchingReplays)
    {
        _allReplays.Replays()[_replayIndex]->Frames()[_replayFrameIndex].Render(_frameBufferA, _activeGame->Collection());
    }
    else
    {
        _activeGame->ActiveFrame().Render(_frameBufferA, _activeGame->Collection());
    }

    _frameTextureA->LoadImageData(_frameBufferA, GLTexture::kNone);
    _frameTextureB->LoadImageData(_frameBufferB, GLTexture::kNone);

    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
    glTexEnvf(GL_TEXTURE_ENV, GL_TEXTURE_ENV_MODE, GL_MODULATE);

    glColor3f(1.0f, 1.0f, 1.0f);
    
    float textureY = (float)Constants::NESHeight / (float)Constants::ScreenTextureHeight;

    const bool splitScreen = _parameters.testingModel;

    if(_parameters.exportingMovie)
    {
        Bitmap correctedFrame(Constants::NESWidth, Constants::NESHeight);
        _frameBufferA.BltTo(correctedFrame, 0, 0);
        correctedFrame.FlipVertical();
        //correctedFrame.FlipBlueAndRed();
        
        correctedFrame.SavePNG("../movie/frame" + Utility::ZeroPadInteger(_movieFrame, 5) + ".png");
        _movieFrame++;
    }

    if(splitScreen)
    {
        _frameTextureA->Bind();
        glBegin(GL_POLYGON);
        glTexCoord2f(0.0f, textureY); glVertex2f(-1.0f, -1.0f);
        glTexCoord2f(1.0f, textureY); glVertex2f( 0.0f, -1.0f);
        glTexCoord2f(1.0f, 0.0f);     glVertex2f( 0.0f,  1.0f);
        glTexCoord2f(0.0f, 0.0f);     glVertex2f(-1.0f,  1.0f);
        glEnd();
        _frameTextureA->UnBind();

        _frameTextureB->Bind();
        glBegin(GL_POLYGON);
        glTexCoord2f(0.0f, textureY); glVertex2f( 0.0f, -1.0f);
        glTexCoord2f(1.0f, textureY); glVertex2f( 1.0f, -1.0f);
        glTexCoord2f(1.0f, 0.0f);     glVertex2f( 1.0f,  1.0f);
        glTexCoord2f(0.0f, 0.0f);     glVertex2f( 0.0f,  1.0f);
        glEnd();
        _frameTextureB->UnBind();
    }
    else
    {
        _frameTextureA->Bind();
        glBegin(GL_POLYGON);
        glTexCoord2f(0.0f, textureY); glVertex2f(-1.0f, -1.0f);
        glTexCoord2f(1.0f, textureY); glVertex2f( 1.0f, -1.0f);
        glTexCoord2f(1.0f, 0.0f);     glVertex2f( 1.0f,  1.0f);
        glTexCoord2f(0.0f, 0.0f);     glVertex2f(-1.0f,  1.0f);
        glEnd();
        _frameTextureA->UnBind();
    }

    mike::GlTextConsoleRender();
}

void Application::Step()
{
    if(_parameters.testingModel)
    {
        
    }
    else if(_parameters.watchingReplays)
    {
        _replayFrameIndex++;
        if(_replayFrameIndex >= int(_allReplays.Replays()[_replayIndex]->Frames().size()))
        {
            _replayFrameIndex = 0;
            _replayIndex++;
            if(_replayIndex >= int(_allReplays.Replays().size())) _replayIndex = 0;
        }
    }
    else if (_parameters.learningOnline)
    {
        // This will call _activeGame->Step()
        _onlineLearningCoordinator->takeStep();
    }
    else
    {
        if(_useAI)
        {
            _controller = _activeGame->SimpleAI();
        }

        if(_parameters.recording)
        {
            if(_replay.Frames().size() < _parameters.framesToRecord)
            {
                _replay.AddFrame(_activeGame->ActiveFrame(), _controller);
                if(_replay.Frames().size() == _parameters.framesToRecord) _replay.SaveToNewReplayFile();
            }
        }

        _activeGame->Step(_controller);
    }
}

void Application::KeyDown(int key)
{
    if(key == 'g')
    {
        if(_parameters.testingModel)
        {
            _testingBaseFrame++;
            const vector<Frame> &curFrames = _allReplays.Replays()[0]->Frames();
            _testingFrames.clear();
            for(int historyFrameIndex = 0; historyFrameIndex < _activeModel.Parameters().historyFrameCount; historyFrameIndex++)
            {
                _testingFrames.push_front(&curFrames[Math::Bound(_testingBaseFrame - historyFrameIndex, 0, int(curFrames.size()) - 1)]);
            }

            for(int lookAheadIndex = 0; lookAheadIndex < _testingLookAhead; lookAheadIndex++)
            {
                _activeModel.AdvanceFrame(Utility::ListToVector(_testingFrames), _predictedFrames[lookAheadIndex]);
                _predictedFrames[lookAheadIndex].controller = curFrames[Math::Bound(_testingBaseFrame + lookAheadIndex + 1, 0, int(curFrames.size()) - 1)].controller;
                _testingFrames.push_front(&_predictedFrames[lookAheadIndex]);
                _testingFrames.pop_back();
            }

            const Frame &baseFrame = _allReplays.Replays()[0]->Frames()[_testingBaseFrame];
            for(UINT variableIndex = 0; variableIndex < baseFrame.observedInts.size(); variableIndex++)
            {
                cout << _activeGame->ObservedIntNames()[variableIndex] << ": " << Utility::IntToString(baseFrame.observedInts[variableIndex]) << endl;
            }
        }
    }
    if(key == 'm')
    {
        _useAI = !_useAI;
        _controller.state = 0;
    }
    _controller.KeyDown(key);
}

void Application::KeyUp(int key)
{
    _controller.KeyUp(key);
}